import pyrealsense2 as rs
import numpy as np


class D435Camera:
    """
    Intel RealSense D435 相机简洁封装。

    统一对外接口：
    - get_align_rgb(): 返回与彩色图对齐的 RGB 图 (np.ndarray, HxWx3, BGR)
    - get_align_depth(): 返回与彩色图对齐的深度图(米) (np.ndarray, HxW, float32)
    - get_align_allframe(): 返回 (rgb, depth) 均为对齐后的图像
    - get_color_intrinsics(): 返回 (fx, fy, cx, cy)
    """

    def __init__(self, width: int = 640, height: int = 480, fps: int = 30, autostart: bool = True) -> None:
        self.width = width
        self.height = height
        self.fps = fps

        self.pipeline = rs.pipeline()
        self.config = rs.config()

        # 配置彩色与深度流
        self.config.enable_stream(rs.stream.color, width, height, rs.format.bgr8, fps)
        self.config.enable_stream(rs.stream.depth, width, height, rs.format.z16, fps)

        self.profile = None
        self.depth_scale = None
        self.align = None
        self._color_intrinsics = None

        if autostart:
            self.start()

    def start(self) -> None:
        self.profile = self.pipeline.start(self.config)

        # 深度尺度（米）
        depth_sensor = self.profile.get_device().first_depth_sensor()
        self.depth_scale = float(depth_sensor.get_depth_scale())

        # 对齐到彩色
        self.align = rs.align(rs.stream.color)

        # 读取彩色相机内参
        color_stream = self.profile.get_stream(rs.stream.color).as_video_stream_profile()
        intr = color_stream.get_intrinsics()
        self._color_intrinsics = (float(intr.fx), float(intr.fy), float(intr.ppx), float(intr.ppy))

    def stop(self) -> None:
        try:
            self.pipeline.stop()
        except Exception:
            pass

    # ----------------------- 基础帧获取 -----------------------
    def _wait_frames(self):
        return self.pipeline.wait_for_frames()

    def get_align_allframe(self):
        """返回 (color_image, depth_in_meters) 对齐后的帧。"""
        frames = self._wait_frames()
        aligned_frames = self.align.process(frames)

        depth_frame = aligned_frames.get_depth_frame()
        color_frame = aligned_frames.get_color_frame()
        if not depth_frame or not color_frame:
            return None, None

        depth_image = np.asanyarray(depth_frame.get_data()).astype(np.float32) * self.depth_scale
        color_image = np.asanyarray(color_frame.get_data())
        return color_image, depth_image

    def get_align_rgb(self):
        frames = self._wait_frames()
        aligned_frames = self.align.process(frames)
        color_frame = aligned_frames.get_color_frame()
        if not color_frame:
            return None
        color_image = np.asanyarray(color_frame.get_data())
        return color_image

    def get_align_depth(self):
        frames = self._wait_frames()
        aligned_frames = self.align.process(frames)
        depth_frame = aligned_frames.get_depth_frame()
        if not depth_frame:
            return None
        depth_image = np.asanyarray(depth_frame.get_data()).astype(np.float32) * self.depth_scale
        return depth_image

    # ----------------------- 其它辅助 -----------------------
    def get_color_intrinsics(self):
        """返回 (fx, fy, cx, cy)。若不可用，返回 None。"""
        return self._color_intrinsics

    # 为统一接口提供别名，便于上层直接调用
    def get_camera_frame(self):  # 与 DetectModel 接口一致
        return self.get_align_rgb()
